from PIL import Image
from math import sqrt
import numpy as np
from tqdm import tqdm
import time
from skimage import color


def legacy_ColourDistance(base_color, img):

    # base_color = np.transpose(base_color, (1, 0))

    # base_color = base_color.reshape((1, *base_color.shape))
    R_1, G_1, B_1 = base_color

    R_2, G_2, B_2 = img[:, :, 0], img[:, :, 1], img[:, :, 2]

    rmean = (R_1 + R_2) / 2.0
    R = R_1 - R_2
    G = G_1 - G_2
    B = B_1 - B_2
    return np.sqrt(
        (2 + rmean / 256) * (R**2) + 4 * (G**2) + (2 + (255 - rmean) / 256) * (B**2)
    )


def ColourDistance(base_color, img):
    """
    计算基础色与图像各像素的CIEDE2000色差
    参数:
        base_color: (R, G, B) 值域[0,255]
        img: (H, W, 3) 值域[0,255]
    返回:
        deltaE_map: (H, W) 色差矩阵
    """
    # 将基础色转为Lab空间
    base_lab = color.rgb2lab(np.array(base_color).reshape(1, 1, 3) / 255.0)
    # base_lab[:, :, 0] = 1.0
    # print(np.min(base_lab), np.max(base_lab))

    # 转换整幅图像到Lab空间
    img_lab = color.rgb2lab(img.astype(np.float32) / 255.0)
    # img_lab[:, :, 0] = 1.0
    # print(np.min(img_lab), np.max(img_lab))
    # print(base_lab.shape, img_lab.shape)
    # delta = (img_lab - base_lab)[:, :, :]
    # delta = delta * delta
    # # print(delta.shape)
    # delta = np.sum(delta, axis=2)
    # return delta

    # 计算CIEDE2000色差
    deltaE_map = color.deltaE_ciede2000(base_lab, img_lab, channel_axis=-1)

    return deltaE_map


def rgb_to_lab(r, g, b):
    # Step 1: Normalize RGB values
    r, g, b = r / 255.0, g / 255.0, b / 255.0

    # Step 2: Convert RGB to XYZ
    x = 0.412453 * r + 0.357580 * g + 0.180423 * b
    y = 0.212671 * r + 0.715160 * g + 0.072169 * b
    z = 0.019334 * r + 0.119193 * g + 0.950227 * b

    # Step 3: Normalize by D65 reference white
    ref_x, ref_y, ref_z = 0.95047, 1.0, 1.08883
    x = x / ref_x
    y = y / ref_y
    z = z / ref_z

    # Step 4: Apply non-linear transformation
    def f(t):
        bigger = t > 0.008856
        t[bigger] = t[bigger] ** (1 / 3)
        t[~bigger] = (t[~bigger] * 7.787) + (16 / 116)
        return t
        # return t ** (1 / 3) if t > 0.008856 else (7.787 * t) + (16 / 116)

    fx = f(x)
    fy = f(y)
    fz = f(z)

    # Step 5: Compute L*, a*, b*
    l = 116 * fy - 16
    a = 500 * (fx - fy)
    b_value = 200 * (fy - fz)

    return l, a, b_value


COLORS = {
    "Black": np.array((0, 0, 0)),
    "Red": np.array((255, 0, 0)),
    "Orange": np.array((255, 165, 0)),
    "Yellow": np.array((255, 255, 0)),
    "Green": np.array((0, 255, 0)),
    "Cyan": np.array((0, 255, 221)),
    "Blue": np.array((0, 0, 255)),
    "White": np.array((255, 255, 255)),
    # "Purple": np.array((255, 0, 255)),
}
color_arrs = np.stack(list(COLORS.values()))
lab_color_arrs = color_arrs
# lab_color_arrs = np.stack(
#     rgb_to_lab(color_arrs[:, 0], color_arrs[:, 1], color_arrs[:, 2])
# )
# lab_color_arrs = np.transpose(lab_color_arrs, (1, 0))
# print(lab_color_arrs)


image_path = "img.jpeg"
image = Image.open(image_path)
image = image.convert("RGB")

# new_img = np.zeros(tuple(reversed((3, *image.size))), dtype=np.uint8)


# for i in tqdm(range(image.size[1])):
#     for j in range(image.size[0]):
#         pixel = image.getpixel((j, i))
#         # print(pixel)
#         distances = []
#         for color, val in COLORS.items():
#             dist = sqrt(
#                 0
#                 + (pixel[0] - val[0]) ** 2
#                 + (pixel[1] - val[1]) ** 2
#                 + (pixel[2] - val[2]) ** 2
#             )
#             distances.append((dist, color))

#         min_dist, color = min(distances)
#         new_img[i, j] = COLORS[color]

# new_im = Image.fromarray(new_img)
# # new_im.show()
# new_im.save("output.png")
start = time.time()
new_img = np.array(image)
# new_img = np.stack(rgb_to_lab(new_img[:, :, 0], new_img[:, :, 1], new_img[:, :, 2]))
# new_img = np.transpose(new_img, (1, 2, 0))
print(new_img.shape)
print(color_arrs.shape)

diffs = np.zeros((*new_img.shape[:2], color_arrs.shape[0]))
legacy_diffs = np.zeros((*new_img.shape[:2], color_arrs.shape[0]))
# for i in range(color_arrs.shape[0]):
#     col_mul = new_img - lab_color_arrs[i]
#     col_mul *= col_mul
#     col_mul = np.sum(col_mul, axis=2)
#     diffs[:, :, i] = col_mul
for i in range(color_arrs.shape[0]):
    ciede = ColourDistance(color_arrs[i], new_img)
    # legacy = legacy_ColourDistance(color_arrs[i], new_img)
    # print(ciede == legacy)
    diffs[:, :, i] = ciede
    # legacy_diffs[:, :, i] = legacy

# print(diffs)
mins = np.argmin(diffs, axis=2)
legacy_mins = np.argmin(legacy_diffs, axis=2)
print(np.sum(mins == legacy_mins) / new_img.shape[0] / new_img.shape[1])

new_img = color_arrs[mins]
# new_img = color_arrs[legacy_mins]
new_im = Image.fromarray(new_img.astype(np.uint8))
new_im.show()
new_im.save("np_output.png")

# new_im_legacy = color_arrs[legacy_mins]
# new_im_legacy = Image.fromarray(new_im_legacy.astype(np.uint8))
# new_im_legacy.show()
# new_im_legacy.save("legacy_output.png")

print(f"耗时：{time.time() - start}s")
